#!/usr/bin/env python3

# Copyright (c) 2005-2014 Fpemud <fpemud@sina.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import os
import re
import glob
import pathlib
import collections
from ._util import Util
from ._handy import Handy
from ._po import Arch
from ._po import KernelType
from ._exception import BootEntryInvalid


class BootEntrySpec(collections.abc.Mapping):

    def __init__(self, data):
        if "arch" not in data:
            assert False
        if data["arch"] not in [x.value for x in list(Arch)]:
            assert False

        if "kernel_type" not in data:
            assert False
        if data["kernel_type"] not in [x.value for x in list(KernelType)]:
            assert False

        if "kernel_params" not in data:
            assert False
        if not isinstance(data["kernel_params"], dict):
            assert False

        if "kernel_addon_params" not in data:
            assert False
        for params in data["kernel_addon_params"].values():
            if not isinstance(params, dict):
                assert False

        if "initramfs_params" not in data:
            assert False
        if not isinstance(data["initramfs_params"], dict):
            assert False

        self._data = data

    @property
    def arch(self):
        return Arch(self._data["arch"])

    @property
    def kernel_type(self):
        return KernelType(self._data["kernel_type"])

    @property
    def kernel_addons(self):
        return self._data["kernel_addon_params"].keys()

    def get_kernel_params(self):
        return self._data["kernel_params"]

    def get_kernel_addon_params(self, kernel_addon_name):
        return self._data["kernel_addon_params"][kernel_addon_name]

    def get_initramfs_params(self):
        return self._data["initramfs_params"]

    def to_dict(self):
        return self._data

    def __eq__(self, other):
        if type(self) is not type(other):
            return False
        if self._data != other._data:
            return False
        return True

    def __getitem__(self, key):
        return self._data[key]

    def __iter__(self):
        return iter(self._data)

    def __len__(self):
        return len(self._data)


class BootEntry:

    spec_filename_wildcard = "spec-*"
    kernel_filename_wildcard = "kernel-*"
    kernel_config_filename_wildcard = "config-*"
    initramfs_filename_wildcard = "initramfs-*"
    initramfs_tar_filename_wildcard = "initramfs-files-*.tar.bz2"

    def __init__(self, bbki, arch=None, verstr=None, postfix=None, history_entry=False):
        self._bbki = bbki

        if arch is not None:
            assert isinstance(arch, Arch)
            assert Handy.isValidKernelVer(verstr.split("-")[0])
            assert postfix is None
        else:
            assert verstr is None

            # postfix example: x86_64-3.9.11-gentoo-r1
            partList = postfix.split("-")
            assert len(partList) >= 2
            assert Handy.isValidKernelVer(partList[1])

            arch = Arch(partList[0])
            verstr = "-".join(partList[1:])

        if not history_entry:
            self._bootDir = self._bbki._fsLayout.get_boot_dir()
        else:
            self._bootDir = self._bbki._fsLayout.get_boot_history_dir()

        # get spec file
        specFile = os.path.join(self._bootDir, "spec-" + arch.value + "-" + verstr)
        if not os.path.exists(specFile):
            raise BootEntryInvalid("boot entry spec file %s not found" % (specFile))

        # parse spec file
        self._verstr = None
        self._spec = None
        self._kernelVer = None
        self._addonVerDict = None
        if True:
            kernelParams = None
            addonParamsDict = None
            initramfsParams = None
            kernelVer = None
            addonVerDict = None

            lineList = pathlib.Path(specFile).read_text().split("\n")

            def __parseParameters(myName, s):
                ret = {}
                if s != "":
                    for item in s.split(" "):
                        tlist = item.split("=")
                        if len(tlist) == 1:
                            ret[tlist[0]] = None
                        elif len(tlist) == 2:
                            ret[tlist[0]] = tlist[1]
                        else:
                            raise BootEntryInvalid("invalid %s parameter" % (myName))
                return ret

            for line in lineList:
                m = re.fullmatch("kernel-params: (.*)", line)
                if m is not None:
                    if kernelParams is not None:
                        raise BootEntryInvalid("duplicate \"kernel-params\" found in %s" % (specFile))
                    kernelParams = __parseParameters("kernel", m.group(1))

            addonParamsDict = {}
            for line in lineList:
                m = re.fullmatch("kernel-addon-(\\S+)-params: (.*)", line)
                if m is not None:
                    if m.group(1) in addonParamsDict:
                        raise BootEntryInvalid("duplicate \"kernel-addons-%s-params\" found in %s" % (m.group(1), specFile))
                    addonParamsDict[m.group(1)] = __parseParameters("kernel addon", m.group(2))

            for line in lineList:
                m = re.fullmatch("initramfs-params: (.*)", line)
                if m is not None:
                    if initramfsParams is not None:
                        raise BootEntryInvalid("duplicate \"initramfs-params\" found in %s" % (specFile))
                    initramfsParams = __parseParameters("initramfs", m.group(1))

            for line in lineList:
                m = re.fullmatch("kernel-version: (.*)", line)
                if m is not None:
                    if kernelVer is not None:
                        raise BootEntryInvalid("duplicate \"kernel-version\" found in %s" % (specFile))
                    kernelVer = m.group(1)

            addonVerDict = {}
            for line in lineList:
                m = re.fullmatch("kernel-addon-(\\S+)-version: (.*)", line)
                if m is not None:
                    if m.group(1) in addonVerDict:
                        raise BootEntryInvalid("duplicate \"kernel-addon-%s-version\" found in %s" % (m.group(1), specFile))
                    addonVerDict[m.group(1)] = m.group(2)

            if kernelParams is None:
                raise BootEntryInvalid("no \"kernel-params\" found in %s" % (specFile))

            if addonParamsDict.keys() != addonVerDict.keys():
                raise BootEntryInvalid("inconsistent \"kernel-addons-params\" and \"kernel-addons-version\" found in %s" % (specFile))

            if initramfsParams is None:
                raise BootEntryInvalid("no \"initramfs-params\" found in %s" % (specFile))

            if kernelVer is None:
                raise BootEntryInvalid("no \"kernel-version\" found in %s" % (specFile))

            if verstr != kernelVer and not verstr.startswith(kernelVer + "-"):
                raise BootEntryInvalid("invalid \"kernel-version\" in %s" % (specFile))

            self._verstr = verstr
            self._spec = BootEntrySpec({
                "arch": arch.value,
                "kernel_type": KernelType.LINUX.value,
                "kernel_params": kernelParams,
                "kernel_addon_params": addonParamsDict,
                "initramfs_params": initramfsParams,
            })
            self._kernelVer = kernelVer
            self._addonVerDict = addonVerDict

    @property
    def spec(self):
        return self._spec

    @property
    def postfix(self):
        # string, eg: "x86_64-3.9.11-gentoo-r1"
        return self._spec.arch.value + "-" + self._verstr

    @property
    def verstr(self):
        # string, eg: "3.9.11-gentoo-r1"
        return self._verstr

    @property
    def spec_filename(self):
        # string, eg: "spec-x86_64-3.9.11-gentoo-r1"
        return "spec-" + self.postfix

    @property
    def spec_filepath(self):
        # string, eg: "/boot/spec-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.spec_filename)

    @property
    def kernel_filename(self):
        # string, eg: "kernel-x86_64-3.9.11-gentoo-r1"
        return "kernel-" + self.postfix

    @property
    def kernel_filepath(self):
        # string, eg: "/boot/kernel-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.kernel_filename)

    @property
    def kernel_config_filename(self):
        # string, eg: "config-x86_64-3.9.11-gentoo-r1"
        return "config-" + self.postfix

    @property
    def kernel_config_filepath(self):
        # string, eg: "/boot/config-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.kernel_config_filename)

    @property
    def kernel_modules_dirpath(self):
        # string, eg: "/lib/modules/5.1.14-gentoo-r1"
        return self._bbki._fsLayout.get_kernel_modules_dir(self._verstr)

    @property
    def firmware_dirpath(self):
        # string, eg: "/lib/firmware"
        return self._bbki._fsLayout.get_firmware_dir()

    @property
    def initramfs_filename(self):
        # string, eg: "initramfs-x86_64-3.9.11-gentoo-r1"
        return "initramfs-" + self.postfix

    @property
    def initramfs_filepath(self):
        # string, eg: "/boot/initramfs-x86_64-3.9.11-gentoo-r1"
        return os.path.join(self._bootDir, self.initramfs_filename)

    @property
    def initramfs_tar_filename(self):
        # string, eg: "initramfs-files-x86_64-3.9.11-gentoo-r1.tar.bz2"
        return "initramfs-files-" + self.postfix + ".tar.bz2"

    @property
    def initramfs_tar_filepath(self):
        # string, eg: "/boot/initramfs-x86_64-3.9.11-gentoo-r1.tar.bz2"
        return os.path.join(self._bootDir, self.initramfs_tar_filename)

    def is_historical(self):
        return (self._bootDir == self._bbki._fsLayout.get_boot_history_dir())

    def get_kernel_version(self):
        # string, eg: "3.9.11"
        return self._kernelVer

    def get_kernel_addon_version(self, kernel_addon_name):
        return self._addonVerDict[kernel_addon_name]

    def has_kernel_files(self):
        if not os.path.exists(self.spec_filepath):
            return False
        if not os.path.exists(self.kernel_filepath):
            return False
        if not os.path.exists(self.kernel_config_filepath):
            return False
        if not os.path.exists(self.kernel_modules_dirpath):
            return False
        if not os.path.exists(self.firmware_dirpath):
            return False
        return True

    def has_initramfs_files(self):
        if not os.path.exists(self.initramfs_filepath):
            return False
        if not os.path.exists(self.initramfs_tar_filepath):
            return False
        return True

    # FIXME
    def __eq__(self, other):
        assert False


class BootEntryUtils:

    @staticmethod
    def get_filepaths(bootEntry):
        return [
            bootEntry.spec_filepath,
            bootEntry.kernel_filepath,
            bootEntry.kernel_config_filepath,
            bootEntry.initramfs_filepath,
            bootEntry.initramfs_tar_filepath,
        ]

    @staticmethod
    def get_firmware_filenames(bootEntry):
        ret = set()

        for fullKoFn in glob.glob(os.path.join(bootEntry.kernel_modules_dirpath, "**", "*.ko"), recursive=True):
            ret |= set(Util.kmodGetModuleInfoByFile(fullKoFn)["firmwares"])

        fwExtFileRecordFn = os.path.join(bootEntry.firmware_dirpath, "firmware.extra-files")
        if os.path.exists(fwExtFileRecordFn):
            ret |= set(Util.readListFile(fwExtFileRecordFn))
            ret.add("firmware.extra-files")

        return sorted(list(ret))

    @classmethod
    def get_firmware_filepaths(cls, bootEntry):
        return [os.path.join(bootEntry.firmware_dirpath, x) for x in cls.get_firmware_filenames(bootEntry)]

    @staticmethod
    def hashBootEntry(bootEntry):
        return hash((bootEntry._bootDir, bootEntry._spec.arch, bootEntry._verstr))

    @staticmethod
    def compareBootEntries(bootEntry1, bootEntry2):
        assert isinstance(bootEntry1, BootEntry) and isinstance(bootEntry2, BootEntry)
        assert bootEntry1._bbki is bootEntry2._bbki

        if bootEntry1._bootDir != bootEntry2._bootDir:
            return False
        if bootEntry1._spec.arch != bootEntry2._spec.arch:
            return False
        if bootEntry1._verstr != bootEntry2._verstr:
            return False
        return True

    @classmethod
    def moveToHistory(cls, bootEntry):
        Util.forceMkdir(bootEntry._bbki._fsLayout.get_boot_history_dir())
        for fullfn in cls.get_filepaths(bootEntry):
            if os.path.exists(fullfn):
                Util.forceMoveToDir(fullfn, bootEntry._bbki._fsLayout.get_boot_history_dir())

    @staticmethod
    def getRedundantKernelModulesDirs(bootEntryList):
        bbki = bootEntryList[0]._bbki

        if not os.path.exists(bbki._fsLayout.get_kernel_modules_dir()):
            return []

        ret = glob.glob(os.path.join(bbki._fsLayout.get_kernel_modules_dir(), "*"))  # mark /lib/modules/* (no recursion) as to-be-deleted
        for be in bootEntryList:
            try:
                ret.remove(bbki._fsLayout.get_kernel_modules_dir(be.verstr))         # don't delete files belongs to a boot-entry
            except ValueError:
                pass
        return ret

    @classmethod
    def getRedundantFirmwareFiles(cls, bootEntryList):
        bbki = bootEntryList[0]._bbki

        if not os.path.exists(bbki._fsLayout.get_firmware_dir()):
            return []

        tset = set(Util.globDirRecursively(bbki._fsLayout.get_firmware_dir()))   # mark /lib/firmware/* (recursive) as to-be-deleted
        for be in bootEntryList:
            for fp in cls.get_firmware_filepaths(be):
                if True:
                    # don't delete this file (and its ancestor directories), it belongs to a boot-entry
                    fullfn = fp
                    while fullfn != "/":
                        tset.discard(fullfn)
                        fullfn = os.path.dirname(fullfn)
                if os.path.islink(fp):
                    # don't delete the file pointed by this symlink (and its ancestor directories)
                    fullfn = os.path.normpath(os.path.join(os.path.dirname(fp), os.readlink(fp)))
                    while fullfn != "/":
                        tset.discard(fullfn)
                        fullfn = os.path.dirname(fullfn)

        return sorted(list(tset))
